#pragma once
#include <assert.h>
#include <functional>
#include <iostream>
#include <utility>

using namespace std;

template <class K, class V> struct AVLTreeNode {
  pair<K, V> _kv;
  AVLTreeNode<K, V> *left;
  AVLTreeNode<K, V> *right;
  AVLTreeNode<K, V> *parent;
  // 平衡因子
  int _bf;
  AVLTreeNode(const pair<K, V> &kv)
      : left(nullptr), right(nullptr), parent(nullptr), _kv(kv), _bf(0) {}
};

template <class K, class V> class AVLTree {
  typedef AVLTreeNode<K, V> Node;

public:
  pair<Node *, bool> Insert(const pair<K, V> &kv) {
    if (_root == nullptr) {
      _root = new Node(kv);
      return make_pair(_root, true);
    }
    Node *cur = _root;
    Node *parent = nullptr;
    while (cur) {
      if (cur->_kv.first > kv.first) {
        cur = cur->left;
        parent = cur;
      } else if (cur->_kv.first < kv.first) {
        cur = cur->right;
        parent = cur;
      } else {
        return make_pair(cur, false);
      }
    }
    cur = new Node(kv);
    Node *temp = cur;
    if (parent->_kv.first > kv.first) {
      parent->left = cur;
    } else {
      parent->right = cur;
    }
    cur->parent = parent;
    while (cur != _root) {
      if (cur == parent->left) {
        parent->_bf--;
      } else if (cur == parent->right) {
        parent->_bf++;
      }
      if (parent->_bf == 0) {
        break;
      } else if (parent->_bf == -1 || parent->_bf == 1) {
        cur = parent;
        parent = parent->parent;
      } else if (parent->_bf == -2 || parent->_bf == 2) {
        if (parent->_bf == -2) {
          if (cur->_bf == -1) {
            RoateR(parent);
          } else {
            RoateLR(parent);
          }
        } else if (parent->_bf == 2) {
          if (cur->_bf == 1) {
            RoateL(parent);
          } else {
            RoateRL(parent);
          }
        }
        break;
      } else {
        assert(false);
      }
      return make_pair(temp, true);
    }
  }
  // 左旋
  void RoateL(Node *parent) {
    Node *subR = parent->right;
    Node *subRL = subR->left;
    parent->right = subRL;
    if (subRL) {
      subRL->parent = parent;
    }
    Node *pParent = parent->parent;
    subR->left = parent;
    parent->parent = subR;
    if (!pParent) {
      _root = subR;
      subR->parent = nullptr;
    } else {
      if (parent == pParent->left) {
        pParent->left = subR;
      } else {
        pParent->right = subR;
      }
      subR->parent = pParent;
    }
    parent->_bf = 0;
    subR->_bf = 0;
  }
  // 右旋
  void RoateR(Node *parent) {
    Node *subL = parent->left;
    Node *subLR = subL->right;
    parent->left = subLR;
    if (subLR) {
      subLR->parent = parent;
    }
    Node *pParent = parent->parent;
    subL->right = parent;
    parent->parent = subL;
    if (!pParent) {
      _root = subL;
      subL->parent = nullptr;
    } else {
      if (parent == pParent->left) {
        pParent->left = subL;
      } else {
        pParent->right = subL;
      }
      subL->parent = pParent;
    }
    subL->_bf = subLR->_bf = 0;
  }
  // 左右双旋
  void RoateLR(Node *parent) {
    Node *subL = parent->left;
    Node *subLR = parent->right;
    int bf = subLR->_bf;
    RoateL(parent->left);
    RoateR(parent);
    if (bf == -1) {
      subL->_bf = 0;
      subLR->_bf = 0;
      parent->_bf = 1;
    } else if (bf == 1) {
      subL->_bf = -1;
      subLR->_bf = 0;
      parent->_bf = 0;
    } else if (bf == 0) {
      subL->_bf = 0;
      subLR->_bf = 0;
      parent->_bf = 0;
    } else {
      assert(false);
    }
  }
  // 右左双旋
  void RoateRL(Node *parent) {
    Node *subR = parent->right;
    Node *subRL = subR->left;
    int bf = subRL->_bf;
    RoateR(parent->right);
    RoateL(parent);
    if (bf == -1) {
      subRL->_bf = 0;
      subR->bf = 1;
      parent->bf = 0;
    } else if (bf == 1) {
      subRL->_bf = 0;
      subR->bf = 0;
      parent->bf = -1;
    } else if (bf == 0) {
      subRL->_bf = 0;
      subR->bf = 0;
      parent->bf = 0;
    } else {
      assert(false);
    }
  }
  void Inorder() { _InOrder(_root); }
  bool IsAVLTree() { return _IsBalanceTree(_root); }
  Node *Find(const K &key) {
    Node *cur = _root;
    while (cur) {
      if (key < cur->_kv.first) {
        cur = cur->left;
      } else if (key > cur->_kv.first) {
        cur = cur->right;
      } else {
        return cur;
      }
    }
    return nullptr;
  }
  bool Modify(const K &key, const V &value) {
    Node *pos = Find(key);
    if (pos == nullptr) {
      return false;
    }
    pos->_kv.second = value;
    return true;
  }
  V &operator[](const K &key) {
    pair<Node *, bool> pos = Insert(make_pair(key, V()));
    Node *node = pos.first;
    return node->_kv.second;
  }
  bool Erase(const K &key) {
    Node *parent = nullptr;
    Node *cur = _root;
    Node *deleteParentNode = nullptr;
    Node *deleteNode = nullptr;
    while (cur) {
      if (key > cur->_kv.first) {
        parent = cur;
        cur = cur->right;
      } else if (key < cur->_kv.first) {
        parent = cur;
        cur = cur->left;
      } else {
        if (cur->left == nullptr) // 左子树为空
        {
          if (cur == _root) {
            _root = _root->right;
            if (_root) {
              _root->parent = nullptr;
            }
            delete cur;
            return true;
          } else {
            deleteParentNode = parent;
            deleteNode = cur;
          }
          break;
        } else if (cur->right == nullptr) // 右子树为空
        {
          if (cur == _root) {
            _root = _root->left;
            if (_root) {
              _root->parent = nullptr;
            }
            delete cur;
            return true;
          } else {
            deleteParentNode = parent;
            deleteNode = cur;
          }
          break;
        } else // 左右子树都不是空，要进行替换删除
        {
          // 找到我们删除节点的右子树中最小的作为提换
          Node *minParent = cur;
          Node *minRight = cur->right;
          while (minRight->left) {
            minParent = minRight;
            minRight = minRight->left;
          }
          // 替换里面的内容
          cur->_kv.first = minRight->_kv.first;
          cur->_kv.second = minRight->_kv.second;
          // 替换删除的节点
          deleteParentNode = minParent;
          deleteNode = minRight;
          break;
        }
      }
      if (deleteParentNode == nullptr) {
        return false;
      }
      // 记录节点
      Node* dNode = deleteNode;
      Node* dpNode = deleteParentNode;
      // 更新平衡因子
      while (deleteNode != _root) {
        if (deleteNode == deleteParentNode->left) {
          deleteParentNode->_bf++;
        } else if (deleteParentNode == deleteParentNode->right) {
          deleteParentNode->_bf--;
        }
        if (deleteParentNode->_bf == 0) {
          deleteNode = deleteParentNode;
          deleteParentNode = deleteParentNode->parent;
        } else if (deleteParentNode->_bf == -1 || deleteParentNode->_bf == 1) {
          break;
        } else if (deleteParentNode->_bf == -2 || deleteParentNode->_bf == 2) {
          if (deleteParentNode->_bf == -2) {
            if (deleteParentNode->left->_bf == -1) {
              Node *temp = deleteParentNode->left;
              RoateR(deleteParentNode);
              deleteParentNode = temp;
            } else if (deleteParentNode->left->_bf == 1) {
              Node *temp = deleteParentNode->left->right;
              RoateLR(deleteParentNode);
              deleteParentNode = temp;
            } else {
              Node *temp = deleteParentNode->left;
              RoateR(deleteParentNode);
              deleteParentNode->_bf = 1;
              deleteParentNode->right->_bf = -1;
              break;
            }
          } else {
            if (delParentPos->_right->_bf == -1) {
              Node *tmp = delParentPos->_right->_left; 
              RotateRL(delParentPos); 
              delParentPos = tmp;     
            } else if (delParentPos->_right->_bf == 1) {
              Node *tmp = delParentPos->_right; 
              RotateL(delParentPos); 
              delParentPos = tmp;   
            } else                   
            {
              Node *tmp = delParentPos->_right; 
              RotateL(delParentPos);
              delParentPos = tmp;    
              delParentPos->_bf = -1;
              delParentPos->_left->_bf = 1;
              break; 
            }
          }
          deleteNode = deleteParentNode;
          deleteParentNode = deleteParentNode->parent;
        }else {
          assert(false);
        }
      }
      // 进行删除
      if(dNode->left == nullptr){
        if(dNode == dpNode->left) {
          dpNode->left = dNode->right;
          if(dNode->right) {
            dNode->right->parent = dpNode;
          }
        }else {
          dpNode->right = dNode->right;
          if(dNode->right) {
            dNode->right->parent = dpNode;
          }
        }
      }else {
        if(dNode == dpNode->left) {
          dpNode->left = dNode->left;
          if(dNode->left) {
            dNode->left->parent = dpNode;
          }
        }else {
          dpNode->right = dNode->left;
          if(dNode->left) {
            dNode->left->parent = dpNode;
          }
        }
      }
      delete dNode;
      return true;
    }
  private:
    // 中序遍历
    void _InOrder(Node * root) {
      if (root == nullptr) {
        return;
      }
      _InOrder(root->left);
      cout << root->_kv.first << ":" << root->_kv.second << endl;
      _InOrder(root->right);
    }
    // 求高度
    int _Height(Node * root) {
      if (root == nullptr) {
        return 0;
      }
      return max(_Height(root->left) + 1, _Height(root->right) + 1);
    }
    // 判断平衡性
    bool _IsBalanceTree(Node * root) {
      // 空树也是AVL树
      if (root == nullptr) {
        return true;
      }
      int leftHeight = _Height(root->left);
      int rightHeight = _Height(root->right);
      int diff = rightHeight - leftHeight;
      if (abs(diff) >= 2) {
        cout << "这节点的高度差异常：" << root->_kv.first << endl;
        return false;
      }
      if (root->_bf != diff) {
        cout << "这个节点的平衡因子异常：" << root->_kv.first << endl;
        return false;
      }
      return _IsBalanceTree(root->left) && _IsBalanceTree(root->right);
    }
    private:
      Node *_root;
};
